# code: utf-8
import copy
import time

import numpy as np

try:
    import DetailDesignV2.jxq.group as group
except:
    import group as group
import random
import itertools as it
from numpy import linalg as LA
try:
    import DetailDesignV2.jxq.model_operator as model_operator
except:
    import model_operator as model_operator

# 调整RPS点的顺序，与距离最近的clamp点靠近
def modify_rps_order(obj,rps_coor):
    if len(rps_coor)>1:
        clamp_points=get_clamp_point(obj)
        clamp_points_list=[v[0] for v in clamp_points]
        rps_combination_dict=[]
        rps_dist_list=[]
        for temp_rps in it.permutations(rps_coor,len(rps_coor)):
            dis_rps_clamp=0.0
            for clamp_point,rps_point in zip(clamp_points_list,temp_rps):
                dis_rps_clamp+=LA.norm(clamp_point-rps_point)
            rps_dist_list.append(dis_rps_clamp)
            rps_combination_dict.append(temp_rps)
        min_index = rps_dist_list.index(min(rps_dist_list))
        rps_coor=rps_combination_dict[min_index]
    print("调整后的rps坐标为：", rps_coor)
    return rps_coor


# 获取夹具点的坐标
def get_clamp_point(clamp=None):
    point = []
    clamp_point=clamp.obj.get('clamp_points', [])
    for i in range(len(clamp_point)):
        for j in range(len(clamp_point[i])):

            point.append(clamp.get_point(clamp_point[i][j]))
    # print(point)
    points = np.array(point).reshape(len(clamp_point), len(clamp_point[0]), 3)
    # print("夹紧点的坐标为：", points)
    return points




def get_a_new_example(obj):
    temp_obj=copy.deepcopy(obj)
    para_dict=copy.deepcopy(obj.obj['args'])
    new_para_dict=copy.deepcopy(para_dict)
    for key in para_dict:
        new_para_dict[key]['v']=key['v']+random.random()
    temp_obj.obj['args']=new_para_dict
    new_obj = group.Body(temp_obj.obj)
    new_obj = model_operator.update_model_by_history(new_obj, obj)
    return new_obj

def get_a_batch_of_examples(obj,n):
    result=[]
    for i in range(n):
        new_obj=get_a_new_example(obj)
        result.append(new_obj)
    return result


# # 调整RPS点的顺序，与距离最近的clamp点靠近
# def rps_modify(obj, rps_coor):
#     new_rps = copy.deepcopy(rps_coor)
#     if len(rps_coor) > 1:
#         clamp_points = get_clamp_point(obj)
#         for i in range(len(rps_coor)):
#             min_distance = []
#             for j in range(len(clamp_points)):
#                 dis_rps_clamp = 0.0
#                 for rps_x, clamp_x in zip(rps_coor[i], clamp_points[j][0]):
#                     dis_rps_clamp += (rps_x - clamp_x) ** 2
#                 min_distance.append(np.sqrt(dis_rps_clamp))
#             min_index = min_distance.index(min(min_distance))
#             new_rps[min_index] = copy.deepcopy(rps_coor[i])
#     rps_coor = new_rps
#     print("调整后的rps坐标为：", rps_coor)
#     return rps_coor




# # 获得各个RPS点与夹具点之间的距离总和
# def get_diff_1(obj, rps_coor):
#     total_loss = 0
#
#     if len(rps_coor) == 1:
#         clamp_points = get_clamp_point(obj)
#         for i in range(clamp_points.shape[1]):
#             dis_rps_clamp = 0
#             for rps_x, clamp_x in zip(rps_coor[0], clamp_points[0][i]):
#                 dis_rps_clamp += (rps_x - clamp_x) ** 2
#             total_loss += np.sqrt(dis_rps_clamp)
#     else:
#         clamp_points = get_clamp_point(obj)
#         # print("更新后的夹具点坐标为：", self.clamp_points)
#         for i in range(len(rps_coor)):
#             for j in range(len(clamp_points[i])):
#                 dis_rps_clamp = 0
#                 for rps_x, clamp_x in zip(rps_coor[i], clamp_points[i][j]):
#                     dis_rps_clamp += (rps_x - clamp_x) ** 2
#                 total_loss += np.sqrt(dis_rps_clamp)
#     # print("夹具点与RPS点的距离：", total_loss)
#     return total_loss


def get_rps_grad(value):
    obj, rps_coor, global_dict, base_loss, dx, grad_dict, ele=value
    obj_temp = group.Body(obj.obj, global_cache=global_dict)
    obj_temp = model_operator.update_model_by_history(obj_temp, obj)

    after_loss = get_rps_diff(obj_temp, rps_coor)
    print(after_loss)

    # obj_temp.show()
    grad_temp = (after_loss - base_loss) / dx
    if abs(grad_temp) < 1e-6:
        grad_dict[ele['name']] = 0
    else:
        if abs(base_loss / grad_temp) < 5:
            grad_dict[ele['name']] = base_loss / grad_temp
        else:
            grad_dict[ele['name']] = grad_temp

    return grad_temp

# 获得各个RPS点与夹具点之间的距离总和
def get_rps_diff(obj, rps_coor):
    total_loss = 0
    clamp_points=get_clamp_point(obj)
    for clamp_point_list,rps_point in zip(clamp_points,rps_coor):
        new_rps_list=[rps_point for i in range(len(clamp_point_list))]
        for clamp_point, rps_point in zip(clamp_point_list, new_rps_list):
            total_loss += LA.norm(clamp_point - rps_point)
    return total_loss

# 夹具各部件尺寸更新
def update_args(obj, rps_coor, lr=0.05, dx=0.1, iters=40):
    multiprocess_state=False
    if multiprocess_state:
        import multiprocessing
        try:
            from part import Part, PART_DIR, GROUP_DIR,Global_Cache
        except:
            from .part import Part, PART_DIR, GROUP_DIR,Global_Cache
        multiprocessing.freeze_support()  # Windows 平台要加上这句，避免 RuntimeError
        mgr = multiprocessing.Manager()
        global_dict=mgr.dict()
        for k,v in Global_Cache.Cache.items():
            global_dict[k]=v
    else:
        global_dict=None

    start=time.time()


    if multiprocess_state:
        grad_dict=mgr.dict()
    else:
        grad_dict = {}
    if multiprocess_state:
        pool = multiprocessing.Pool()
    for i in range(iters):
        base_loss = get_rps_diff(obj, rps_coor)
        print(base_loss)
        args_list=[]

        for ele in obj.obj['args']:
            if 'isconst' not in ele or not ele['isconst']:
                ele['v'] += dx
                if multiprocess_state:
                    args_list.append((copy.deepcopy(obj),rps_coor,global_dict,base_loss,dx,grad_dict,ele))
                else:
                    get_rps_grad((obj,rps_coor,global_dict,base_loss,dx,grad_dict,ele))

                # obj_temp = group.Body(obj.obj,global_cache=global_dict)
                # obj_temp=model_operator.update_model_by_history(obj_temp,obj)
                #
                # after_loss = get_rps_diff(obj_temp, rps_coor)
                #
                # # obj_temp.show()
                # grad_temp = (after_loss - base_loss) / dx
                # if abs(grad_temp)<1e-6:
                #     grad_dict[ele['name']]=0
                # else:
                #     if abs(base_loss/grad_temp)<5:
                #         grad_dict[ele['name']] =base_loss/ grad_temp
                #     else:
                #         grad_dict[ele['name']]=grad_temp
                ele['v'] -= dx
        if multiprocess_state:
            # pool.apply_async(get_rps_grad, args=(obj, rps_coor, global_dict, base_loss, dx, grad_dict, ele))
            pool.map(get_rps_grad,args_list)

        for ele in obj.obj['args']:
            if ele['name'] in grad_dict:
                ele['v'] = ele['v'] - lr * grad_dict[ele['name']]
                if 'range' in ele:
                    if ele['v']<ele['range'][0]:
                        ele['v']=ele['range'][0]
                    if ele['v']>ele['range'][1]:
                        ele['v']=ele['range'][1]
        # dt=time.time()-start_time
        # print(dt)
        # print(grad_dict)
        new_obj = group.Body(obj.obj)
        new_obj=model_operator.update_model_by_history(new_obj, obj)
        # new_obj.show()
        obj=new_obj
        # if get_rps_diff(obj, rps_coor)<0.5:
        #     break


        # print("after {} iters: the loss is :{}".format(i, get_diff(obj, rps_coor)))
        # print("迭代{}后，clamp_coor坐标为：{}".format(i, get_clamp_point(obj)))
    # print("迭代{}的时间为：{}".format(iters, (time.time() - start_time)))
    # obj.show()
    dt=time.time()-start
    print(dt)
    return obj



          
# # 通过平移更新夹具空间位置
# def update_trans(obj, rps_coor, lr=1, dx=1):
#     base_loss = get_diff(obj, rps_coor)
#
#     obj=model_operator.move_clamp_direction(obj,direction=[dx, 0, 0])
#     x_loss = get_diff(obj, rps_coor)
#     x_grad = - lr * (x_loss - base_loss) / dx
#     obj=model_operator.move_clamp_direction(obj,direction=[-dx, 0, 0])
#
#     obj=model_operator.move_clamp_direction(obj,direction=[0, dx, 0])
#     y_loss = get_diff(obj, rps_coor)
#     y_grad = - lr * (y_loss - base_loss) / dx
#     obj=model_operator.move_clamp_direction(obj,direction=[0, -dx, 0])
#
#     obj=model_operator.move_clamp_direction(obj,direction=[0, 0, dx])
#     z_loss = get_diff(obj, rps_coor)
#     z_grad = - lr * (z_loss - base_loss) / dx
#     obj=model_operator.move_clamp_direction(obj,direction=[0, 0, -dx])
#
#     obj=model_operator.move_clamp_direction(obj,direction=[x_grad, y_grad, z_grad])
    # print("三个方向的移动梯度为{}, {}, {}".format(x_grad, y_grad, z_grad))


    
# 模型更新
def model_update(obj, rps_coor, env=None, section=None, air_point=None, lr=1, dx=1, iters=1):

    # 调整RPS点的顺序，与距离最近的clamp点靠近
    rps_coor=modify_rps_order(obj,rps_coor)
    # rps_coor = rps_modify(obj, rps_coor)
    print("调整后的rps_coor：", rps_coor)
    # for i in range(20):
    #     obj=get_rapid_to_target(obj,rps_coor)

    obj = update_args(obj, rps_coor, lr=lr, dx=dx, iters=iters)
    print("更新后夹具点的坐标为：{}，距离为：{}".format(get_clamp_point(obj), get_rps_diff(obj, rps_coor)))
    print("更新后的零件尺寸为：", obj.args)
    # print("模型更新时间为：", time.time() - start_time)
    return obj




if __name__ == "__main__":
    rps1 = [[-370.5, 0, -95]]
    #
    # rps2 = [
    #         [-370.5, 80, -95],
    #         [-370.5, -80, -95],
    #     ]
    #
    # rps3 = [
    #
    #         [-185.0,  0, -90],
    #         [-360.5, 80, -95],
    #         [-360.5, -80, -95],
    #     ]
    # # Env = trimesh.load("work_2_xy.stl")
    # x = group.Group("clamp")
    # # x.show()
    # section = [1, 0, 0]
    # air_point = [0, 0, 0]
    # model_update(obj=x, rps_coor=rps1, env=None, section=None, air_point=None, lr=1, dx=0.5, iters=100)




